from flask import Flask, request, jsonify
from flask_cors import CORS
import base64
import cv2
import torch
import os
import tempfile
import urllib.request
import numpy as np

app = Flask(__name__)
CORS(app)  # 允许所有来源的请求

# 动态设备分配
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def load_model(model_path):
    """
    加载保存为完整模型的 .pth 文件
    """
    model = torch.load(model_path, map_location=device)
    model.eval()  # 切换到评估模式
    print(f"Model successfully loaded on {device}")
    return model

def preprocess_image(image_path):
    """
    预处理3D NumPy数组，转为 5D 张量输入 (batch_size, channels, depth, height, width)
    """
    image_array = np.load(image_path)  # 加载3D图像数据
    image_array = (image_array - np.min(image_array)) / (np.max(image_array) - np.min(image_array))  # 归一化到 [0, 1]

    # 增加维度形成 (1, channels, depth, height, width)
    if len(image_array.shape) == 2:  # 如果是2D图像数据，添加通道维度和批次维度
        image_array = np.expand_dims(image_array, axis=0)  # (1, height, width)
        image_array = np.expand_dims(image_array, axis=0)  # (1, 1, height, width)
    else:  # 如果是3D图像数据，添加批次维度
        image_array = np.expand_dims(image_array, axis=0)  # (1, depth, height, width)

    # 添加通道维度
    image_array = np.expand_dims(image_array, axis=1)  # (1, 1, depth, height, width)

    # 确保通道数为16
    if image_array.shape[2] != 16:
        image_array = np.repeat(image_array, 16, axis=2)  # (1, 1, 16, height, width)

    # 转为 PyTorch 张量并移动到设备
    tensor = torch.tensor(image_array, dtype=torch.float32).to(device)

    # 调试：打印张量形状和数据范围
    print(f"Preprocessed tensor shape: {tensor.shape}")
    print(f"Tensor min: {tensor.min()}, max: {tensor.max()}")
    print(f"Tensor mean: {tensor.mean()}, std: {tensor.std()}")
    return tensor  # 返回 PyTorch 张量

def apply_highlight(input_slice, prediction_slice, threshold=0.5):
    """
    在预测结果上应用荧光效果
    input_slice: 输入图像的切片（灰度图）
    prediction_slice: 预测结果的切片
    threshold: 阈值，用于确定目标区域
    """
    # 将预测结果转换为二值图像
    _, binary_prediction = cv2.threshold(prediction_slice, threshold, 255, cv2.THRESH_BINARY)

    # 将灰度图像转换为 RGB 图像
    input_rgb = cv2.cvtColor(input_slice, cv2.COLOR_GRAY2RGB)

    # 创建一个与输入图像相同大小的全零 RGB 图像
    highlighted = np.zeros_like(input_rgb)

    # 将二值图像中的白色区域设置为绿色
    highlighted[binary_prediction == 255] = [0, 255, 0]

    # 将高亮图像与输入图像合并
    result = cv2.addWeighted(input_rgb, 0.7, highlighted, 0.3, 0)
    return result

@app.route('/predict', methods=['POST'])
def predict_image():
    if 'filePath' not in request.form:
        return jsonify({'error': 'No file path provided'})

    file_path = request.form['filePath']
    print(f"Received file path: {file_path}")



    try:
        with urllib.request.urlopen(file_path) as response:
            image_data = response.read()
            print(f"Downloaded image data length: {len(image_data)}")

            # 使用 cv2.imdecode 从内存中读取图片
            image = cv2.imdecode(np.frombuffer(image_data, np.uint8), cv2.IMREAD_GRAYSCALE)
            if image is None:
                return jsonify({'error': 'Failed to load image'})

            # 将PNG文件转换为NumPy数组并保存为.npy文件
            image_array = np.array(image)
            temp_npy_path = tempfile.NamedTemporaryFile(delete=False, suffix='.npy').name
            np.save(temp_npy_path, image_array)
            print(f"Saved image array to: {temp_npy_path}")

            # 加载模型
            model_path = "./models/best_model.pth"  # 替换为实际路径
            model = load_model(model_path)

            # 加载测试数据
            input_tensor = preprocess_image(temp_npy_path)  # 返回 torch.Tensor

            # 可视化输入图像的中间切片
            depth = input_tensor.shape[2] // 2  # 计算深度中间位置

            # 模型预测
            with torch.no_grad():  # 禁用梯度计算
                output = model(input_tensor)
                output = torch.sigmoid(output)  # 转换为概率图
                output = output[0, 1].detach().cpu().numpy()  # 提取预测结果 (3D NumPy array)

            # 将输入张量转换为 NumPy 数组
            input_array = input_tensor[0, 0].cpu().numpy()  # (depth, height, width)

            # 提取中间切片
            input_slice = input_array[depth, :, :]
            prediction_slice = output[depth, :, :]

            # 应用荧光效果
            highlighted_image = apply_highlight(input_slice, prediction_slice)

            # 将高亮图像转换为 uint8 类型并缩放到 [0, 255]
            highlighted_image = (highlighted_image * 255).astype(np.uint8)

            # 将图像编码为Base64
            _, buffer = cv2.imencode('.png', highlighted_image)
            base64_image = base64.b64encode(buffer).decode('utf-8')

            # 删除临时文件
            os.remove(temp_npy_path)

            return jsonify({'predictedImage': base64_image})
    except Exception as e:
        print(f"Error: {e}")
        return jsonify({'error': 'Failed to process image'})

if __name__ == '__main__':
    app.run(host='localhost', port=5000)
